#!/usr/bin/env python
"""ZMQClientNode: ROS node logic for sending images to a ZMQ inference server.

This module contains the `ZMQClientNode` class. It intentionally keeps ROS/ZMQ
integration here while using helpers from `transforms` and `utils` modules.
"""
from __future__ import annotations
import rospy
from sensor_msgs.msg import Image
from std_msgs.msg import Empty
from cv_bridge import CvBridge
import zmq
import json
import numpy as np
import cv2
import threading
import time

from zmqclient.transforms import transform_pose_camera_to_ros, euler_rad_to_deg
from zmqclient.utils import encode_png
from arm_control.msg import PosCmd
import math


class ZMQClientNode:
    def __init__(self):
        self.server_addr = rospy.get_param('~server_addr', 'tcp://192.168.8.26:5555')
        self.color_topic = rospy.get_param('~color_topic', '/camera/color/image_raw')
        self.depth_topic = rospy.get_param('~depth_topic', '/camera/depth/image_raw')
        self.timeout = rospy.get_param('~timeout', 10.0)

        self.bridge = CvBridge()
        self.synced_color = None
        self.synced_depth = None
        self.lock = threading.Lock()

        # ZMQ setup
        self.ctx = zmq.Context()
        self.socket = self.ctx.socket(zmq.REQ)
        self.socket.RCVTIMEO = int(self.timeout * 1000)
        self.socket.connect(self.server_addr)

        rospy.Subscriber(self.color_topic, Image, self.color_callback, queue_size=1)
        rospy.Subscriber(self.depth_topic, Image, self.depth_callback, queue_size=1)

        self.trigger_topic = rospy.get_param('~trigger_topic', '/grasp_trigger')
        rospy.Subscriber(self.trigger_topic, Empty, self.trigger_callback, queue_size=1)

        rospy.loginfo(f'ZMQ client connecting to {self.server_addr}')

        self.continuous = rospy.get_param('~continuous', False)
        if self.continuous:
            self.worker = threading.Thread(target=self._worker_loop, daemon=True)
            self.worker.start()

        # Arm control parameters and ROS pub/sub for commanding the arm
        self.arm_pub_rate = rospy.get_param('~arm_pub_rate', 30)
        self.use_scurve = rospy.get_param('~use_scurve', True)
        self.tighten_duration = rospy.get_param('~tighten_duration', 0.2)
        self.stage_duration = rospy.get_param('~stage_duration', 0.5)
        self.final_approach_duration = rospy.get_param('~final_approach_duration', 1.0)
        self.wait_current_pose_timeout = rospy.get_param('~wait_current_pose_timeout', 2.0)
        self.publish_traj_debug = rospy.get_param('~publish_traj_debug', False)

        # Publisher for arm commands and subscriber for current end-effector state
        self.arm_pub = rospy.Publisher('/arm_cmd', PosCmd, queue_size=10)
        self.arm_state_sub = rospy.Subscriber('/arm_status_ee', PosCmd, self._arm_state_callback)

        # current pose storage (populated from /arm_status_ee)
        self.current_pose = {
            'x': 0.0, 'y': 0.0, 'z': 0.0,
            'roll': 0.0, 'pitch': 0.0, 'yaw': 0.0,
            'gripper': 0.0
        }
        self.current_pose_received = False

    def color_callback(self, msg: Image):
        try:
            cv = self.bridge.imgmsg_to_cv2(msg, desired_encoding='rgb8')
        except Exception:
            try:
                cv = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
                cv = cv[:, :, ::-1]
            except Exception:
                return
        with self.lock:
            self.synced_color = cv.copy()

    def depth_callback(self, msg: Image):
        try:
            d = self.bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
        except Exception:
            return
        with self.lock:
            self.synced_depth = d.copy()

    def trigger_callback(self, msg: Empty):
        rospy.loginfo('Received grasp trigger message, requesting one inference')
        with self.lock:
            color = None if self.synced_color is None else self.synced_color.copy()
            depth = None if self.synced_depth is None else self.synced_depth.copy()

        if color is None or depth is None:
            rospy.logwarn('No image available at trigger time; ignoring trigger')
            return

        resp = self.request_once(color, depth)
        if resp is None:
            rospy.logerr('Request failed or timed out')
            return
        if not resp.get('ok', False):
            rospy.loginfo('No valid grasp received on trigger: %s' % resp.get('reason', ''))
            return
        self.execute_grasp(resp)
    
    # ------------------ Arm helper methods ------------------
    def _arm_state_callback(self, msg: PosCmd):
        """Callback to receive current end-effector pose/status from `/arm_status_ee`.
        We prefer explicit roll/pitch/yaw fields on the incoming message.
        """
        try:
            self.current_pose['x'] = float(msg.x)
            self.current_pose['y'] = float(msg.y)
            self.current_pose['z'] = float(msg.z)
        except Exception:
            pass

        # Prefer roll/pitch/yaw fields
        try:
            self.current_pose['roll'] = float(msg.roll)
            self.current_pose['pitch'] = float(msg.pitch)
            self.current_pose['yaw'] = float(msg.yaw)
        except Exception:
            pass

        try:
            self.current_pose['gripper'] = float(msg.gripper)
        except Exception:
            pass

        self.current_pose_received = True

    def request_once(self, color, depth):
        color_b = encode_png(color, is_depth=False)
        depth_b = encode_png(depth, is_depth=True)

        if color_b is None or depth_b is None:
            rospy.logwarn('Failed to encode images for one-shot request')
            return None

        header = {'ts': time.time()}
        rospy.loginfo('Sending one-shot request to %s (timeout=%ds)' % (self.server_addr, int(self.timeout)))
        t0 = time.time()
        try:
            self.socket.send_multipart([json.dumps(header).encode('utf-8'), color_b, depth_b])
            rospy.loginfo('Request sent (t=%.3fs)' % (time.time() - t0))
        except Exception as e:
            rospy.logerr(f'ZMQ send error (one-shot): {e} -- resetting socket')
            try:
                self.socket.close(0)
            except Exception:
                pass
            self.socket = self.ctx.socket(zmq.REQ)
            self.socket.RCVTIMEO = int(self.timeout * 1000)
            try:
                self.socket.connect(self.server_addr)
            except Exception:
                pass
            return None

        try:
            rospy.loginfo('Waiting for reply...')
            reply = self.socket.recv_string()
            rospy.loginfo('Reply received (t=%.3fs)' % (time.time() - t0))
        except Exception as e:
            rospy.logerr(f'ZMQ recv error or timeout (one-shot): {e} -- resetting socket')
            try:
                self.socket.close(0)
            except Exception:
                pass
            self.socket = self.ctx.socket(zmq.REQ)
            self.socket.RCVTIMEO = int(self.timeout * 1000)
            try:
                self.socket.connect(self.server_addr)
            except Exception:
                pass
            return None

        try:
            resp = json.loads(reply)
        except Exception as e:
            rospy.logerr(f'Invalid JSON reply (one-shot): {e} -- {reply}')
            return None

        return resp

    def _worker_loop(self):
        rate = rospy.Rate(5)
        while not rospy.is_shutdown():
            with self.lock:
                color = None if self.synced_color is None else self.synced_color.copy()
                depth = None if self.synced_depth is None else self.synced_depth.copy()

            if color is None or depth is None:
                rate.sleep()
                continue

            color_b = encode_png(color, is_depth=False)
            depth_b = encode_png(depth, is_depth=True)

            if color_b is None or depth_b is None:
                rospy.logwarn('Failed to encode images, skipping send')
                rate.sleep()
                continue

            header = {'ts': time.time()}
            try:
                self.socket.send_multipart([json.dumps(header).encode('utf-8'), color_b, depth_b])
            except Exception as e:
                rospy.logerr(f'ZMQ send error: {e} -- resetting socket')
                try:
                    self.socket.close(0)
                except Exception:
                    pass
                self.socket = self.ctx.socket(zmq.REQ)
                self.socket.RCVTIMEO = int(self.timeout * 1000)
                try:
                    self.socket.connect(self.server_addr)
                except Exception:
                    pass
                rate.sleep()
                continue

            try:
                reply = self.socket.recv_string()
            except Exception as e:
                rospy.logerr(f'ZMQ recv error or timeout: {e} -- resetting socket')
                try:
                    self.socket.close(0)
                except Exception:
                    pass
                self.socket = self.ctx.socket(zmq.REQ)
                self.socket.RCVTIMEO = int(self.timeout * 1000)
                try:
                    self.socket.connect(self.server_addr)
                except Exception:
                    pass
                rate.sleep()
                continue

            try:
                resp = json.loads(reply)
            except Exception as e:
                rospy.logerr(f'Invalid JSON reply: {e} -- {reply}')
                rate.sleep()
                continue

            if not resp.get('ok', False):
                rospy.loginfo('No valid grasp received: %s' % resp.get('reason', ''))
            else:
                self.execute_grasp(resp)

            if resp.get('shutdown', False):
                rospy.loginfo('Server requested shutdown (debug). Exiting client loop.')
                break

            rate.sleep()
    
    def execute_grasp(self, resp: dict):
        try:
            trans = resp.get('translation')
            rot = resp.get('rotation')
        except Exception:
            rospy.loginfo('Received grasp (malformed resp)')
            return

        try:
            if rot is not None:
                pos_new, quat_xyzw, euler_rad = transform_pose_camera_to_ros(trans, rotation_matrix=rot)
            else:
                q_in = resp.get('quaternion') or resp.get('quat')
                if q_in is None:
                    rospy.logwarn('No rotation info in response; cannot transform pose')
                    return
                pos_new, quat_xyzw, euler_rad = transform_pose_camera_to_ros(trans, quaternion=q_in)
        except Exception as e:
            rospy.logerr(f'Failed to transform pose to ROS frame: {e}')
            return

        pos_str = f"pos=[{pos_new[0]:.4f}, {pos_new[1]:.4f}, {pos_new[2]:.4f}]"

        # human-friendly degrees
        rpy_deg = euler_rad_to_deg(euler_rad)
        if rpy_deg is not None:
            rpy_str = f"rpy(deg)=[{rpy_deg[0]:.2f}, {rpy_deg[1]:.2f}, {rpy_deg[2]:.2f}]"
        else:
            rpy_str = 'rpy=None'

        # also provide original radians explicitly (from transform)
        if euler_rad is not None:
            rpy_rad_str = f"rpy(rad)=[{euler_rad[0]:.4f}, {euler_rad[1]:.4f}, {euler_rad[2]:.4f}]"
        else:
            rpy_rad_str = 'rpy_rad=None'

        # Now apply your requested axis offsets to the degree values:
        # - pitch (index 1) add +90 degrees; if >180 then subtract 180
        # - yaw (index 2) subtract 90 degrees
        adjusted_rpy_deg = None
        adjusted_rpy_rad = None
        if rpy_deg is not None:
            adj = [float(rpy_deg[0]), float(rpy_deg[1]), float(rpy_deg[2])]
            if adj[0] < -90.0:
                adj[0] += 180.0
            if adj[0] > 90.0:
                adj[0] -= 180.0
            adjusted_rpy_deg = (adj[0], adj[1], adj[2])
            # convert adjusted degrees to radians
            adjusted_rpy_rad = tuple(np.radians(adjusted_rpy_deg))

        if adjusted_rpy_deg is not None:
            adj_deg_str = f"adj_rpy(deg)=[{adjusted_rpy_deg[0]:.2f}, {adjusted_rpy_deg[1]:.2f}, {adjusted_rpy_deg[2]:.2f}]"
            adj_rad_str = f"adj_rpy(rad)=[{adjusted_rpy_rad[0]:.4f}, {adjusted_rpy_rad[1]:.4f}, {adjusted_rpy_rad[2]:.4f}]"
        else:
            adj_deg_str = 'adj_rpy_deg=None'
            adj_rad_str = 'adj_rpy_rad=None'

        # quaternion string
        if quat_xyzw is not None:
            quat_str = f"quat=[{quat_xyzw[0]:.4f}, {quat_xyzw[1]:.4f}, {quat_xyzw[2]:.4f}, {quat_xyzw[3]:.4f}]"
        else:
            quat_str = 'quat=None'

        # gripper width scaling: (width / 0.08) * 5
        gw = resp.get('gripper_width')
        gw_scaled = None
        if gw is not None:
            gw_scaled = (float(gw) / 0.08) * 5.0


        rospy.loginfo(
            f"Received grasp: score={resp.get('score')}, {pos_str}, {rpy_str}, {rpy_rad_str}, {adj_deg_str}, {adj_rad_str}, {quat_str}, width={gw}, width_scaled={gw_scaled}"
        )
        # Implement robot motion planning/execution here using pos_new and adjusted_rpy_rad
        target_rpy_rad = adjusted_rpy_rad

        if target_rpy_rad is None:
            rospy.logwarn('No adjusted RPY available; cannot publish roll/pitch/yaw target')
            return

        # Wait for current arm pose (block up to configured timeout)
        t0 = time.time()
        while not self.current_pose_received and (time.time() - t0) < float(self.wait_current_pose_timeout) and not rospy.is_shutdown():
            rospy.sleep(0.01)

        if not self.current_pose_received:
            rospy.logwarn(f'No current arm pose received within {self.wait_current_pose_timeout}s; will send one-shot target')

        # Build start and end arrays in the order: x,y,z, roll, pitch, yaw, gripper
        start = np.array([
            float(self.current_pose.get('x', 0.0) - 0.04),
            float(self.current_pose.get('y', 0.0)),
            float(self.current_pose.get('z', 0.0) + 0.05),
            float(self.current_pose.get('roll', 0.0)),
            float(self.current_pose.get('pitch', 0.0)),
            float(self.current_pose.get('yaw', 0.0)),
            float(self.current_pose.get('gripper', 0.0)),
        ], dtype=float)

        gripper_target = gw_scaled if gw_scaled is not None else float(self.current_pose.get('gripper', 0.0))

        end = np.array([
            float(pos_new[0]),
            float(pos_new[1]),
            float(pos_new[2]),
            float(target_rpy_rad[0]),
            float(target_rpy_rad[1]),
            float(target_rpy_rad[2]),
            float(gripper_target),
        ], dtype=float)

        rospy.loginfo(f'Publishing arm trajectory to target: pos={end[0]:.3f},{end[1]:.3f},{end[2]:.3f} rpy(rad)={end[3]:.3f},{end[4]:.3f},{end[5]:.3f} gripper={end[6]:.3f}')

        try:
            self._execute_three_stage_sequence(start, end)
        except Exception as e:
            rospy.logerr(f'Error while executing arm sequence: {e}')

    def _s_curve(self, t: float) -> float:
        """Simple s-curve smoothing for t in [0,1]."""
        if t < 0.5:
            return 2.0 * t * t
        else:
            return 1.0 - 2.0 * (1.0 - t) * (1.0 - t)

    def _publish_trajectory(self, start: np.ndarray, end: np.ndarray, duration: float):
        """Publish linear (or s-curve) interpolated trajectory between start and end.
        start/end arrays are [x,y,z,roll,pitch,yaw,gripper]."""
        rate_hz = float(self.arm_pub_rate)
        steps = max(1, int(max(1, duration * rate_hz)))
        rate = rospy.Rate(rate_hz)

        for i in range(steps + 1):
            t = float(i) / float(steps)
            s = self._s_curve(t) if self.use_scurve else t
            vals = start + s * (end - start)

            msg = PosCmd()
            # position
            try:
                msg.x = float(vals[0])
                msg.y = float(vals[1])
                msg.z = float(vals[2])
            except Exception:
                pass
            # roll/pitch/yaw
            try:
                msg.roll = float(vals[3])
                msg.pitch = float(vals[4])
                msg.yaw = float(vals[5])
            except Exception:
                pass
            # gripper
            try:
                msg.gripper = float(vals[6])
            except Exception:
                pass

            self.arm_pub.publish(msg)
            if self.publish_traj_debug and (i % max(1, int(steps/5)) == 0 or i == steps):
                rospy.loginfo(f'traj publish: i={i}/{steps} pos={vals[0]:.3f},{vals[1]:.3f},{vals[2]:.3f} rpy={vals[3]:.3f},{vals[4]:.3f},{vals[5]:.3f} gripper={vals[6]:.3f}')
            try:
                rate.sleep()
            except rospy.ROSInterruptException:
                break

    def _execute_three_stage_sequence(self, start: np.ndarray, end: np.ndarray):
        """High-level three-stage move -> tighten -> three-stage return.
        start/end: numpy arrays [x,y,z,roll,pitch,yaw,gripper]
        """
        def is_near_origin(p):
            return (abs(p[0]) < 0.03 and abs(p[1]) < 0.03 and abs(p[2]) < 0.03)

        def is_holding_object(gripper_value):
            return gripper_value > 1.5

        x, y, z = end[0], end[1], end[2]

        # safe points
        safe1 = np.array([0.0, 0.0, 0.1, 0.0, 0.0, 0.0, max(end[6], 5.0)])
        safe2 = np.array([0.2, 0.0, 0.1, end[3], end[4], end[5], max(end[6], 5.0)])
        safe3 = np.array([end[0], end[1], end[2] + 0.2, end[3], end[4], end[5], max(end[6], 5.0)])

        # branch 1
        if is_near_origin(np.array([self.current_pose.get('x',0.0), self.current_pose.get('y',0.0), self.current_pose.get('z',0.0)])) and z < 0.2:
            if x < 0.25 and -0.1 < y < 0.1:
                rospy.logwarn('目标点位于下方碰撞禁区，操作已取消')
                return

            self._publish_trajectory(start, safe1, self.stage_duration)
            self._publish_trajectory(safe1, safe2, self.stage_duration)
            self._publish_trajectory(safe2, safe3, self.final_approach_duration)
            self._publish_trajectory(safe3, end, self.final_approach_duration)

            self._perform_grip_and_return(end)
            return

        # branch 2
        if (start[2] < 0) and (not is_near_origin(np.array([self.current_pose.get('x',0.0), self.current_pose.get('y',0.0), self.current_pose.get('z',0.0)]))) and is_near_origin(end):
            safe2_back = np.array([0.3, 0.0, 0.1, start[3], start[4], start[5], start[6]])
            safe1_back = safe1
            self._publish_trajectory(start, safe2_back, self.stage_duration)
            self._publish_trajectory(safe2_back, safe1_back, self.stage_duration)
            if is_holding_object(self.current_pose.get('gripper', 0.0)):
                final_pos = np.array([0.0, 0.0, 0.2, 0.0, 0.0, 0.0, self.current_pose.get('gripper', 0.0)])
                self._publish_trajectory(safe1_back, final_pos, self.stage_duration)
            else:
                self._publish_trajectory(safe1_back, end, self.stage_duration)
            return

        # branch 3
        if is_near_origin(end):
            if is_holding_object(self.current_pose.get('gripper', 0.0)):
                safe_pos = np.array([0.0, 0.0, 0.2, 0.0, 0.0, 0.0, self.current_pose.get('gripper', 0.0)])
                self._publish_trajectory(start, safe_pos, self.stage_duration)
            else:
                self._publish_trajectory(start, end, self.stage_duration)
            return

        # default
        self._publish_trajectory(start, end, self.stage_duration)
        self._perform_grip_and_return(end)

    def _perform_grip_and_return(self, target_pose: np.ndarray):
        """Tighten gripper and perform three-stage return to safe location.
        target_pose is [x,y,z,roll,pitch,yaw,gripper]
        """
        x, y, z = target_pose[0], target_pose[1], target_pose[2]
        q_roll, q_pitch, q_yaw = target_pose[3], target_pose[4], target_pose[5]
        gripper = float(target_pose[6])

        current = np.array([
            float(self.current_pose.get('x', 0.0)),
            float(self.current_pose.get('y', 0.0)),
            float(self.current_pose.get('z', 0.0)),
            float(self.current_pose.get('roll', 0.0)),
            float(self.current_pose.get('pitch', 0.0)),
            float(self.current_pose.get('yaw', 0.0)),
            float(self.current_pose.get('gripper', 0.0)),
        ], dtype=float)

        # 1. tighten gripper
        tightened_gripper = max(2.0, gripper - 1.8)
        tightened_pos = np.array([x, y, z, q_roll, q_pitch, q_yaw, tightened_gripper])
        self._publish_trajectory(current, tightened_pos, self.tighten_duration)

        # 2. three-stage return
        tightened = tightened_pos
        stage1_pos = np.array([x, y, 0.12, q_roll, q_pitch, q_yaw, tightened_gripper])
        stage2_pos = np.array([0.0, 0.0, 0.15, 0.0, 0.0, 0.0, tightened_gripper])
        final_pos = np.array([0.0, 0.0, 0.1, 0.0, 0.0, 0.0, tightened_gripper])

        self._publish_trajectory(tightened, stage1_pos, self.stage_duration)
        self._publish_trajectory(stage1_pos, stage2_pos, self.stage_duration)
        self._publish_trajectory(stage2_pos, final_pos, self.stage_duration)
